#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>
#include <string>

using namespace std;

struct Node {
    string ch;
    int freq;
    Node* left;
    Node* right;

    Node(string ch, int freq) : ch(ch), freq(freq), left(nullptr), right(nullptr) {}// 构造函数
};

// 比较节点频率的小顶堆
struct Compare {
    bool operator()(Node* left, Node* right) {
        return left->freq > right->freq;
    }
};

//构建哈夫曼树
/**
 * @brief 构建哈夫曼树
 *
 * 根据给定的字符频率映射，构建哈夫曼树。
 *
 * @param freqMap 字符频率映射，其中键为字符，值为对应的频率
 *
 * @return 构建的哈夫曼树的根节点指针
 */
Node* buildHuffmanTree(const std::unordered_map<string, int>& freqMap) {
    // 最小堆
    std::priority_queue<Node*, std::vector<Node*>, Compare> minHeap;

    // 遍历频率映射表，将每个字符及其频率创建为节点，并加入最小堆
    for (auto pair : freqMap) {
        minHeap.push(new Node(pair.first, pair.second));// pair.first char ，pair.second int freq
    }

    // 当最小堆中节点数量不等于1时，继续合并节点
    while (minHeap.size() != 1) {
        // 取出频率最小的节点作为左子节点
        Node* left = minHeap.top();
        minHeap.pop();

        // 取出频率次小的节点作为右子节点
        Node* right = minHeap.top();
        minHeap.pop();

        // 合并左右子节点，创建一个新的节点，其频率为左右子节点频率之和
        Node* merged = new Node("", left->freq + right->freq);
        // 将左子节点指向新创建的节点的左子节点
        merged->left = left;
        // 将右子节点指向新创建的节点的右子节点
        merged->right = right;

        // 将合并后的节点加入最小堆
        minHeap.push(merged);
    }

    // 返回最小堆中唯一的节点，即哈夫曼树的根节点
    return minHeap.top();
}
//生成哈夫曼编码
/**
 * @brief 生成哈夫曼编码
 *
 * 根据给定的哈夫曼树（Huffman Tree）节点和字符串，生成哈夫曼编码，并将编码存储在给定的哈希映射中。
 *
 * @param root 哈夫曼树的根节点
 * @param str 当前路径字符串
 * @param huffmanCode 存储哈夫曼编码的哈希映射
 */
void generateCodes(Node* root, const std::string& str, std::unordered_map<string, std::string>& huffmanCode) {
    // 如果根节点为空，则直接返回
    if (root == nullptr)
        return;

    // 如果当前节点没有左子节点和右子节点，表示是叶子节点
    if (!root->left && !root->right) {
        // 将字符和对应的哈夫曼编码存入哈希表中
        huffmanCode[root->ch] = str;
    }

    // 递归处理左子树，并在当前编码后添加'0'
    generateCodes(root->left, str + "0", huffmanCode);
    // 递归处理右子树，并在当前编码后添加'1'
    generateCodes(root->right, str + "1", huffmanCode);
}

string encodeLongestMatch(const string& text, int startIndex, unordered_map<string, string>& huffmanCode) {
    string encodedStr;
    int textLength = text.length();
    
    while (startIndex < textLength) {
        // 尝试匹配最长字符串
        for (int len = textLength - startIndex; len > 0; --len) {
            string substr = text.substr(startIndex, len);
            if (huffmanCode.count(substr)) {
                // 如果找到了匹配的字符串，则使用其哈夫曼编码
                encodedStr += huffmanCode[substr];
                startIndex += len; // 更新起始索引
                break; // 跳出内层循环，继续查找下一个匹配项
            }
        }
        
        // 如果没有找到任何匹配项，则编码单个字符（作为回退策略）
        if (startIndex < textLength) {
            encodedStr += huffmanCode[string(1, text[startIndex++])];
        }
    }
    
    return encodedStr;
}

//编码
/**
 * @brief 编码字符串
 *
 * 根据给定的哈夫曼编码表，将输入的字符串进行编码。
 *
 * @param text 待编码的字符串
 * @param huffmanCode 哈夫曼编码表，字符到编码字符串的映射
 *
 * @return 编码后的字符串
 */
std::string encode(const std::string& text, std::unordered_map<string, std::string>& huffmanCode) {
    // std::string encodedStr;
    // for (char ch : text) {
    //     encodedStr += huffmanCode[string {ch}];
    // }
    // return encodedStr;
    return encodeLongestMatch(text, 0, huffmanCode);
}
//解码
/**
 * @brief 解码二叉树表示的字符串
 *
 * 根据给定的二叉树和编码字符串，解码出原始字符串。
 *
 * @param root 二叉树的根节点指针
 * @param encodedStr 编码后的字符串
 *
 * @return 解码后的原始字符串
 */
std::string decode(Node* root, const std::string& encodedStr) {
    std::string decodedStr;
    Node* current = root;
    for (char bit : encodedStr) {//循环遍历字符串
        if (bit == '0') {
            current = current->left;
        } else {
            current = current->right;
        }

        if (!current->left && !current->right) {//如果当前节点是叶子节点，则将其字符添加到解码字符串中，并重置当前节点为根节点，开始新一轮的搜索，直到搜索到叶子节点为止。
            decodedStr += current->ch;
            current = root;
        }
    }
    return decodedStr;
}


int main() {
    std::string text = "this is an example for huffman encoding";

    // 计算每个字符的频率
    std::unordered_map<string, int> freqMap;
    for (char ch : text) {
        freqMap[string{ch}]++;
    }
    
    freqMap[string {"this"}] = 10; //指定一个字符串的频率
    //TODO
    /*
        自定义一些长的字符串
    */
    
    // 构建哈夫曼树
    Node* root = buildHuffmanTree(freqMap);

    // 生成哈夫曼编码表
    std::unordered_map<string, std::string> huffmanCode;
    generateCodes(root, "", huffmanCode);

    // 输出哈夫曼编码表
    std::cout << "Huffman Codes:\n";
    for (auto pair : huffmanCode) {
        std::cout << pair.first << ": " << pair.second << '\n';
    }

    // 编码
    std::string encodedStr = encode(text, huffmanCode);
    std::cout << "\nEncoded String:\n" << encodedStr << '\n';

    // 解码
    std::string decodedStr = decode(root, encodedStr);
    std::cout << "\nDecoded String:\n" << decodedStr << '\n';

    return 0;
}
 